import time

from langchain.callbacks import StreamingStdOutCallbackHandler
from typing import Generator


class QianMaiStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
    # def __init__(self):
    tokens = []
    # 记得结束后这里置true
    finish = False

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.tokens.append(token)

    def on_llm_end(self, response, **kwargs) -> None:
        self.finish = True

    def on_llm_error(self, error: Exception, **kwargs) -> None:
        self.tokens.append(str(error))

    # def generate_tokens(self) -> Generator:
    #     while not self.finish:  # or self.tokens:
    #         if self.tokens:
    #             yield self.tokens.pop(0)
    #         else:
    #             time.sleep(1)